package com.sunday.common.mq.rocket.brave.study.base.e10_Brave_Trace;

import brave.Span;
import brave.Tracer;
import brave.Tracing;
import brave.messaging.MessagingTracing;
import brave.propagation.Propagation;
import brave.propagation.TraceContext;
import brave.propagation.TraceContextOrSamplingFlags;
import lombok.extern.slf4j.Slf4j;
import org.apache.rocketmq.client.producer.SendResult;
import org.apache.rocketmq.common.message.Message;
import org.apache.rocketmq.common.message.MessageExt;

import static brave.Span.Kind.CONSUMER;
import static brave.Span.Kind.PRODUCER;

@Slf4j
public final class RocketTracing {

    static final String
            ROCKET_TOPIC = "rocket.topic",
            ROCKET_TAGS = "rocket.tags",
            ROCKET_MSG_ID = "rocket.msgId";

    final Tracing tracing;
    final Tracer tracer;
    final MessagingTracing messagingTracing;
    final String[] traceIdHeaders;
    final TraceContext.Extractor<MessageProducerRequest> producerExtractor;
    final TraceContext.Extractor<MessageConsumerRequest> consumerExtractor;
    final TraceContext.Injector<MessageProducerRequest> producerInjector;
    final TraceContext.Injector<MessageConsumerRequest> consumerInjector;

    public RocketTracing(Tracing tracing) {
        this.tracing = tracing;
        this.tracer = tracing.tracer();
        this.messagingTracing = MessagingTracing.create(tracing);
        Propagation<String> propagation = messagingTracing.propagation();
        this.traceIdHeaders = propagation.keys().toArray(new String[0]);
        this.producerExtractor = propagation.extractor(MessageProducerRequest.GETTER);
        this.consumerExtractor = propagation.extractor(MessageConsumerRequest.GETTER);
        this.producerInjector = propagation.injector(MessageProducerRequest.SETTER);
        this.consumerInjector = propagation.injector(MessageConsumerRequest.SETTER);
    }

    /**
     * 生产者消息前置处理
     */
    public <T> T postProcessMessage(Message message, CallbackHandler<Message, T> handler) throws Exception {

        MessageProducerRequest request = new MessageProducerRequest(message);
        Span span = tracer.nextSpan().name(request.operation()).start();
        //将链路信息注入到userProperties中
        producerInjector.inject(span.context(), request);
        /**
         * span.kind是分布式追踪系统中的一个标记，用于指示跟踪操作的类型。它通常用于OpenTracing和OpenTelemetry等跟踪系统中。
         * span.kind有两个可能的取值：
         * client：表示跟踪操作是作为客户端发起的请求，通常用于标识发送远程调用的操作。例如，一个HTTP请求、RPC调用或发送消息到消息队列等。
         * server：表示跟踪操作是作为服务端接收和处理请求的操作。通常用于标识处理来自客户端的请求的操作。
         * 通过使用span.kind标记，可以在分布式系统中准确地追踪和记录请求的路径，并帮助开发人员理解请求的方向和流向。这对于性能分析、故障排查和监控系统行为非常有用。
         */
        span.kind(PRODUCER).name("publish");
        maybeTag(span, ROCKET_TOPIC, message.getTopic());
        maybeTag(span, ROCKET_TAGS, message.getTags());
        //将span放在“范围”中，这样日志记录器等下游代码就可以看到跟踪id
//        Tracer.SpanInScope scope = tracer.withSpanInScope(span);
//        try (scope) {
        Throwable error = null;
        try {
            T res = handler.get(message);
            if (res != null && res instanceof SendResult result) {
                String msgId = result.getMsgId();
                maybeTag(span, ROCKET_MSG_ID, msgId);
            }
            log.info("[Response][Ack] {}", res);
            return res;
        } catch (Throwable t) {
            error = t;
            log.error("[Response][Nack][ex] cause : [{}]", t.toString(), t);
            throw t;
        } finally {
            //除非您处理异常，否则您可能不知道操作失败!
            if (error != null) span.error(error); //报错处理
            //始终完成跨度
            //注意，范围独立于跨度。一定要完成一段时间。
            span.finish();
        }
    }

    /**
     * 消费者消息处理
     */
    public void invoke(MessageExt message, ExecuteHandler<MessageExt> handler) {
        MessageConsumerRequest request = new MessageConsumerRequest(message);
        TraceContextOrSamplingFlags traceContextOrSamplingFlags = consumerExtractor.extract(request);
//        Tracer tracer = tracing.tracer();
        Span consumerSpan;
        if (traceContextOrSamplingFlags != null && traceContextOrSamplingFlags.context() != null) {
            consumerSpan = tracer.nextSpan(traceContextOrSamplingFlags);
//                consumerSpan = tracer.toSpan(traceContextOrSamplingFlags.context());
        } else {
            consumerSpan = tracer.newTrace().name(request.operation());
        }
        Span listenerSpan = tracer.newChild(consumerSpan.context());
        // isNoop() 方法用于检查一个 Span 是否处于 "noop" 模式。 "noop" 模式是指一种特殊的跟踪模式，其中跟踪器不会执行任何实际的操作或记录任何信息。
        if (!consumerSpan.isNoop()) {
            setConsumerSpan(consumerSpan, message);
            // incur timestamp overhead only once
            long timestamp = tracing.clock(consumerSpan.context()).currentTimeMicroseconds();
            consumerSpan.start(timestamp);
            long consumerFinish = timestamp + 1L; // save a clock reading
            consumerSpan.finish(consumerFinish);
            // not using scoped span as we want to start with a pre-configured time
            listenerSpan.name("on-message").start(consumerFinish);
        }
        Throwable error = null;
        try (Tracer.SpanInScope scope = tracer.withSpanInScope(consumerSpan)) {
            for (String traceIDHeader : traceIdHeaders) message.getProperties().remove(traceIDHeader);
            handler.doWith(message);
        } catch (Throwable t) {
            error = t;
            log.error("[Response][Nack][ex] cause : [{}]", t.toString(), t);
            throw t;
        } finally {
            if (error != null) listenerSpan.error(error);
            listenerSpan.finish();
        }
    }

    void setConsumerSpan(Span span, MessageExt message) {
        span.name("next-message").kind(CONSUMER);
        maybeTag(span, ROCKET_TOPIC, message.getTopic());
        maybeTag(span, ROCKET_TAGS, message.getTags());
        maybeTag(span, ROCKET_MSG_ID, message.getMsgId());
    }

    void maybeTag(Span span, String tag, String value) {
        if (value != null) span.tag(tag, value);
    }

    @FunctionalInterface
    public interface ExecuteHandler<T> {
        void doWith(T t);
    }

    @FunctionalInterface
    public interface CallbackHandler<T, R> {
        R get(T t) throws Exception;
    }

}
